# link: https://mp.weixin.qq.com/s/NTct2_AYhz4Z8q5MYtBQcA

import backtrader as bt
import pandas as pd
import datetime
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


codes = ['sh600276', 'sh600519', 'sh603288']

# 恒瑞医药
data1 = get_data(codes[0])
# 贵州茅台
# data2 = get_data(codes[1])
# # 海天味业
# data3 = get_data(codes[2])

# %%

# 实例化策略
cerebro = bt.Cerebro()

st_date = datetime.datetime(2023, 4, 4)
ed_date = datetime.datetime(2023, 5, 11)

# 添加 600276.SH 的行情数据
datafeed1 = bt.feeds.PandasData(dataname=data1, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed1, name=codes[0])

# 添加 600519.SH 的行情数据
datafeed2 = bt.feeds.PandasData(dataname=data2, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed2, name=codes[1])

# 添加 603288.SH 的行情数据
datafeed3 = bt.feeds.PandasData(dataname=data3, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed3, name=codes[2])


# 第2.3节 新增指标
'''
在回测时，除了常规的高开低收成交量这些行情数据外，还会用到别的指标，
比如选股回测时会用到很多选股因子（PE、PB 、PCF、......），那这些数据又该如何添加进 Backtrader 的数据表格呢？
往 Backtrader 的数据表格里添加指标，就是给数据表格新增列，也就是给数据表格新增 line：
以导入 DataFrame 为例，在继承原始的数据读取类 bt.feeds.PandasData 的基础上，
设置 lines 属性和 params 属性，新的 line 会按其在 lines 属性中的顺序依次添加进数据表格中，
具体对照下面例子的输出部分：
'''


class PandasData_more(bt.feeds.PandasData):
    lines = ('pe', 'pb',)  # 要添加的线
    # 设置 line 在数据源上的列位置
    params = (
        ('pe', -1),
        ('pb', -1),
    )
    # -1表示自动按列名匹配数据，也可以设置为线在数据源中列的位置索引 (('pe',6),('pb',7),)


class TestStrategy(bt.Strategy):
    def __init__(self):
        print("--------- 打印 self.datas 第一个数据表格的 lines ----------")
        print(self.data0.lines.getlinealiases())
        print("pe line:", self.data0.lines.pe)
        print("pb line:", self.data0.lines.pb)


data1['pe'] = 2  # 给原先的data1新增pe指标（简单的取值为2）
data1['pb'] = 3  # 给原先的data1新增pb指标（简单的取值为3）

# 导入的数据 data1 中
cerebro = bt.Cerebro()
st_date = datetime.datetime(2023, 4, 4)
ed_date = datetime.datetime(2023, 5, 11)

# 这里使用上述定义的新类PandasData_more（继承了bt.feeds.PandasData）
datafeed1 = PandasData_more(dataname=data1, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed1, name='600276.SH')
cerebro.addstrategy(TestStrategy)
result = cerebro.run()
